import pandas as pd
import os
import argparse


def setup_mappings():
    """
    根据提供的表格和规则，构建分类映射字典。
    返回一个包含所有映射规则的元组。
    """
    # 一级品类代码 -> 名称
    l1_map = {
        'cy01': '中式正餐', 'cy02': '非中式正餐', 'cy04': '西式快餐', 'cy05': '小吃',
        'cy06': '甜品', 'cy07': '烘焙', 'cy08': '酒吧酒馆', 'cy09': '轻饮',
        'cy10': '咖啡馆', 'cy99': '其他'
    }

    # 二级品类代码 -> 名称
    l2_map = {
        'cy0101': '川湘菜', 'cy0102': '江浙沪菜', 'cy0103': '赣湘鄂菜', 'cy0104': '云贵菜',
        'cy0105': '两广菜', 'cy0106': '西北菜', 'cy0107': '东北菜', 'cy0108': '鲁菜',
        'cy0109': '闽菜', 'cy0110': '自助餐', 'cy0111': '火锅',
        'cy0201': '西餐', 'cy0202': '日式料理', 'cy0203': '韩式料理', 'cy0204': '东南亚菜', 'cy0205': '其他国家正餐',
        'cy0401': '汉堡', 'cy0402': '披萨', 'cy0403': '牛扒/烤肉', 'cy0404': '日式', 'cy0405': '韩式',
        'cy0406': '东南亚', 'cy0407': '轻食/代餐',
        'cy0501': '现制点心', 'cy0502': '炸串', 'cy0503': '麻辣烫', 'cy0504': '烧烤', 'cy0505': '蒸品',
        'cy0506': '凉菜', 'cy0507': '预制黔货', 'cy0508': '其他', 'cy0509': '其他风味',
        'cy0601': '冰淇淋', 'cy0602': '糖水', 'cy0603': '奶酪', 'cy0604': '甜品集合店',
        'cy0701': '综合西式烘焙', 'cy0702': '单一品类西式烘焙', 'cy0703': '中式烘焙',
        'cy0801': '酒吧', 'cy0802': '餐吧',
        'cy0901': '现制饮品', 'cy0902': '非现制饮品', 'cy0903': '茶',
        'cy1001': '咖啡馆', 'cy9999': '其他'
    }

    # 关键词 -> 二级品类代码 (关键词越具体，优先级应该越高，所以放在前面)
    keyword_map = {
        '咖啡': 'cy1001', 'Café': 'cy1001', 'Cà Phê': 'cy1001',
        '英国': 'cy0201', '英式': 'cy0201',
        '越南菜': 'cy0204', '亚洲料理': 'cy0204', '泰国菜': 'cy0204', '印度菜': 'cy0204', '东南亚': 'cy0204',
        '火锅': 'cy0111', 'Lẩu': 'cy0111', '川湘菜': 'cy0101', 'Chilên': 'cy0101', '川菜': 'cy0101', '湘菜': 'cy0101',
        '日式料理': 'cy0202', '日料': 'cy0202', '寿司': 'cy0202', '拉面': 'cy0202', '韩式料理': 'cy0203',
        '韩料': 'cy0203', '韩国': 'cy0203',
        '西餐': 'cy0201', '法餐': 'cy0201', '披萨': 'cy0402', '比萨': 'cy0402', 'Pizza': 'cy0402',
        '牛扒': 'cy0403', '烤肉': 'cy0403', 'BBQ': 'cy0403', '炙': 'cy0403', 'Nướng': 'cy0403',
        '烘焙': 'cy0701', '面包': 'cy0701', '蛋糕': 'cy0701', 'Tiệm bánh': 'cy0701', 'Bánh': 'cy0701',
        '酒吧': 'cy0801', '酒馆': 'cy0801', 'Bar': 'cy0801', 'Pub': 'cy0801', 'Beer': 'cy0801',
    }

    # 小吃二级分类关键词
    snack_keyword_map = {
        '麻辣烫': 'cy0503', '酸辣粉': 'cy0509', '米粉': 'cy0509', '螺蛳粉': 'cy0509', '凉皮': 'cy0506',
        '蒸品': 'cy0505', '包子': 'cy0505', 'Hấp': 'cy0505', '炸串': 'cy0502', '炸鸡': 'cy0502',
        '烧烤': 'cy0504', '凉菜': 'cy0506', '点心': 'cy0501', '小笼包': 'cy0501',
    }
    return l1_map, l2_map, keyword_map, snack_keyword_map


def detect_columns(df: pd.DataFrame) -> dict:
    """
    识别输入表头并返回标准列名映射。
    输出字典的键为标准列名，值为源列名或None。
    标准列名包括：city, primary_category, secondary_category, shop_name,
    avg_price, tel, country, name_intl, email, secondary_category_lv2
    """
    columns = {c: c for c in df.columns}

    def pick(*candidates):
        for c in candidates:
            if c in columns:
                return c
        return None

    mapping = {
        'city': pick('城市', 'city'),
        'primary_category': None,  # 固定“美食”
        'secondary_category': pick('一级品类', 'secondary_category', 'primary_category'),
        'shop_name': pick('品牌名称_intl', 'shop_name', '品牌名称'),
        'avg_price': pick('平均价格', 'avg_price'),
        'tel': pick('电话', 'tel', '电话/手机'),
        'country': pick('country'),  # 如果没有，则由参数/默认填充
        # name_intl 需要优先取英文名：品牌名称 -> name_intl -> 品牌名称_intl
        'name_intl': pick('品牌名称', 'name_intl', '品牌名称_intl'),
        'email': pick('邮箱', 'email'),
        'secondary_category_lv2': pick('二级品类', 'secondary_category_lv2')
    }

    return mapping


def to_target_dataframe(df: pd.DataFrame, country_fallback: str = '') -> pd.DataFrame:
    """
    将输入DataFrame转换为目标列顺序与命名，不引入cy代码，仅保留中文品类名称。
    若缺失country列，则使用country_fallback填充。
    primary_category 固定为"美食"。
    shop_name 与 name_intl 默认均取 品牌名称_intl（若存在），与您的口径保持一致。
    """
    mapping = detect_columns(df)

    def col(src):
        return df[src] if src and src in df.columns else pd.Series([''] * len(df))

    # shop_name 与 name_intl 默认同源：品牌名称_intl → 若不存在则退回 品牌名称
    shop_source = mapping.get('shop_name') or mapping.get('name_intl') or mapping.get('secondary_category')
    name_intl_source = mapping.get('name_intl') or mapping.get('shop_name')

    if not shop_source and '品牌名称' in df.columns:
        shop_source = '品牌名称'
    if not name_intl_source and '品牌名称' in df.columns:
        name_intl_source = '品牌名称'

    out = pd.DataFrame({
        'city': col(mapping.get('city')),
        'primary_category': '美食',
        'secondary_category': col(mapping.get('secondary_category')),
        'shop_name': col(shop_source),
        'avg_price': col(mapping.get('avg_price')),
        'tel': col(mapping.get('tel')),
        'country': col(mapping.get('country')),
        'name_intl': col(name_intl_source),
        'email': col(mapping.get('email')),
        'secondary_category_lv2': col(mapping.get('secondary_category_lv2')),
    })

    if country_fallback:
        out['country'] = out['country'].replace('', country_fallback)

    return out


def map_brand_category(df, l1_map, l2_map, keyword_map, snack_keyword_map):
    """
    处理DataFrame，添加新的分类列。
    """
    l2_to_l1_code = {l2_code: l1_code for l1_code, l2_code_prefix in
                     {'cy01': 'cy01', 'cy02': 'cy02', 'cy04': 'cy04', 'cy05': 'cy05', 'cy06': 'cy06', 'cy07': 'cy07',
                      'cy08': 'cy08', 'cy09': 'cy09', 'cy10': 'cy10', 'cy99': 'cy99'}.items() for l2_code in l2_map if
                     l2_code.startswith(l2_code_prefix)}

    results = []
    for index, row in df.iterrows():
        l1_category_str = str(row.get('类别', '')).strip()
        l2_category_str = str(row.get('品牌类型', '')).strip()

        # 1. 最高优先级规则：检查是否包含'¥'或品类为空
        if '¥' in l1_category_str or (not l1_category_str and not l2_category_str):
            results.append(('cy99', '其他', 'cy9999', '其他'))
            continue

        found_code = None
        search_target_str = ""

        # 2. 决定使用哪个字段进行分类
        if l1_category_str:
            search_target_str = l1_category_str.lower()
        else:  # 如果一级品类为空，则使用二级品类
            search_target_str = l2_category_str.lower()

        # 3. 在选定的字段中进行关键词匹配
        if any(snack_keyword in search_target_str for snack_keyword in ['小吃', 'ăn vặt', 'hè']):
            found_code = 'cy0508'  # 默认为小吃->其他
            for keyword, code in snack_keyword_map.items():
                if keyword in search_target_str:
                    found_code = code;
                    break

        if not found_code:
            for keyword, code in keyword_map.items():
                if keyword.lower() in search_target_str:
                    found_code = code;
                    break

        # 4. 如果在选定字段中未找到任何匹配，则归为"其他"
        if not found_code:
            found_code = 'cy9999'

        # 5. 根据找到的二级品类，派生出一级品类并记录结果
        l2_code = found_code
        l1_code = l2_to_l1_code.get(l2_code, 'cy99')
        l1_name = l1_map.get(l1_code, '其他')
        l2_name = l2_map.get(l2_code, '其他')
        results.append((l1_code, l1_name, l2_code, l2_name))

    new_cols_df = pd.DataFrame(results, index=df.index,
                               columns=['new_一级品类_code', 'new_一级品类_name', 'new_二级品类_code',
                                        'new_二级品类_name'])
    return pd.concat([df, new_cols_df], axis=1)


# --- 主程序入口 ---
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='标准化餐厅数据列名与顺序，不输出cy代码')
    parser.add_argument('--input', required=True, help='输入CSV文件路径')
    parser.add_argument('--output', required=False, help='输出CSV文件路径（可选）。若不提供则自动在输入名后加 _standardized.csv')
    parser.add_argument('--test-limit', type=int, default=0, help='仅处理前N行用于测试，0表示处理全部')
    parser.add_argument('--country', default='', help='当输入没有country列时填充的国家名称，如：印度尼西亚')
    args = parser.parse_args()

    input_filename = args.input
    # 推断输出文件名
    if args.output:
        output_filename = args.output
    else:
        base_name, ext = os.path.splitext(input_filename)
        output_filename = f"{base_name}_standardized.csv"
    test_limit = args.test_limit
    # 智能推断国家：优先命令行 --country，其次从文件名提取（前缀中文国家名）
    if args.country:
        country_fallback = args.country
    else:
        fname = os.path.basename(input_filename)
        if fname.startswith('印度尼西亚'):
            country_fallback = '印度尼西亚'
        elif fname.startswith('马来西亚'):
            country_fallback = '马来西亚'
        elif fname.startswith('菲律宾'):
            country_fallback = '菲律宾'
        elif fname.startswith('越南'):
            country_fallback = '越南'
        elif fname.startswith('哈萨克斯坦'):
            country_fallback = '哈萨克斯坦'
        elif fname.startswith('乌兹别克斯坦'):
            country_fallback = '乌兹别克斯坦'
        else:
            country_fallback = ''

    try:
        df = pd.read_csv(input_filename, encoding='utf-8-sig', dtype=str).fillna('')
        print(f"成功读取文件: {input_filename}")
    except FileNotFoundError:
        print(f"错误：找不到文件 '{input_filename}'。请确认文件名是否正确，以及脚本和CSV文件是否在同一文件夹。")
        exit(1)
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        exit(1)

    if test_limit and test_limit > 0:
        df = df.head(test_limit).copy()
        print(f"测试模式：仅处理前 {len(df)} 行")

    # 若未来需要启用关键词规则，可在此调用 map_brand_category
    # 目前按用户需求：不输出cy代码，仅保留已有的中文一级/二级名称，并统一列名
    standardized_df = to_target_dataframe(df, country_fallback=country_fallback)

    # 输出
    # 确保列顺序完全一致
    target_cols = [
        'city', 'primary_category', 'secondary_category', 'shop_name',
        'avg_price', 'tel', 'country', 'name_intl', 'email', 'secondary_category_lv2'
    ]
    # 对缺失列做兜底
    for col_name in target_cols:
        if col_name not in standardized_df.columns:
            standardized_df[col_name] = ''
    standardized_df = standardized_df[target_cols]

    standardized_df.to_csv(output_filename, index=False, encoding='utf-8-sig')
    print(f"已输出到: {output_filename}，行数: {len(standardized_df)}")